import cv2
import rclpy
from cv_bridge import CvBridge
from rclpy.node import Node
from sensor_msgs.msg import Image


class ImageSubscriber(Node):
    def __init__(self):
        super().__init__('image_subscriber')

        # Create the subscriber. This subscriber will receive an Image from the video_frames topic.
        # The queue size is 0 messages.
        self.subscription = self.create_subscription(Image, 'video_frames', self.listener_callback, 0)

        # Used to suppress unused variable warning
        self.subscription = self.subscription

        # Used to convert between ROS and OpenCV images
        self.br = CvBridge()

    def listener_callback(self, data):
        # Display the message on console
        self.get_logger().info('Receiving video frame')

        # Convert ROS Image message to OpenCV image
        current_frame = self.br.imgmsg_to_cv2(data)

        # Display image
        cv2.imshow("camera", current_frame)
        cv2.waitKey(1)


def main(args=None):
    # Initialize the rclpy library
    rclpy.init(args=args)

    # Create the node
    image_subscriber = ImageSubscriber()

    # Spin the node so the callback function is called
    rclpy.spin(image_subscriber)

    # Destroy the node explicitly (optional - otherwise it will be done automatically when the garbage collector
    # destroys the node object)
    image_subscriber.destroy_node()

    # Shutdown the ROS2 client library for Python
    rclpy.shutdown()


if __name__ == '__main__':
    main()
